"""
试用ELMo模型
"""
import tensorflow
import tensorflow_hub as hub
import tensorflow as tf
from tensorflow.keras import backend as K

r"https://tfhub.dev/google/elmo/3"
r"./elmo_3/saved_model.pb"
r"*****tfhub***/google/elmo/3"
# https://hub.tensorflow.google.cn/google/wiki40b-lm-bg/1


def main():
    # elmo = hub.Module("https://hub.tensorflow.google.cn/google/elmo/3", trainable=True)
    elmo = hub.load("https://hub.tensorflow.google.cn/google/elmo/3").signatures['default']
    # elmo = tf.keras.models.load_model('./elmo_3/')
    embeddings = elmo(
        tensorflow.constant(['the cat is on the mat', 'dogs are in the fog']),
        # signature="default",
        # as_dict=True
    )
    # embeddings = elmo.signatures['default'](['the cat is on the mat', 'dogs are in the fog'])
    print(embeddings)
    print([i for i in embeddings])
    # ['sequence_len', 'word_emb', 'lstm_outputs1', 'default', 'lstm_outputs2', 'elmo']
    print(K.get_value(embeddings['sequence_len']))


if __name__ == '__main__':
    main()
